#!/usr/bin/env python3
"""
2D Damped Harmonic Oscillator Simulation
========================================

This script simulates a 2D damped harmonic oscillator system where:
- The system has equilibrium position at (1, 1)
- Two independent damped oscillators in x and y directions
- Each direction has position and velocity state variables
- System undergoes 1000 simulation steps
- Visualizes the 2D trajectory in phase space
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
import random
import svgwrite
from datetime import datetime
import os

class PeriodicKickDrive:
    def __init__(self, n_kicks=5, velocity_range=2.0, interval_range=(50, 200)):
        """
        Initialize Periodic Kick Drive (PKD) parameters
        
        Parameters:
        - n_kicks: number of kicks in the sequence (N=5)
        - velocity_range: range for random velocity changes
        - interval_range: range for random time intervals between kicks
        """
        self.n_kicks = n_kicks
        self.velocity_range = velocity_range
        self.interval_range = interval_range
        self.kick_sequence = self.generate_kick_sequence()
        self.current_kick_index = 0
        
    def generate_kick_sequence(self):
        """Generate a random PKD sequence"""
        sequence = []
        for i in range(self.n_kicks):
            # Random 2D velocity change
            dvx = random.uniform(-self.velocity_range, self.velocity_range)
            dvy = random.uniform(-self.velocity_range, self.velocity_range)
            velocity_kick = np.array([dvx, dvy])
            
            # Random time interval to next kick
            interval = random.randint(self.interval_range[0], self.interval_range[1])
            
            sequence.append((velocity_kick, interval))
        
        return sequence
    
    def get_current_kick(self):
        """Get current kick and interval, then advance to next"""
        if self.current_kick_index < len(self.kick_sequence):
            kick_data = self.kick_sequence[self.current_kick_index]
            self.current_kick_index += 1
            return kick_data
        else:
            # Reset to beginning for cyclic behavior
            self.current_kick_index = 0
            return self.kick_sequence[self.current_kick_index]
    
    def reset(self):
        """Reset the kick sequence to beginning"""
        self.current_kick_index = 0
        
    def print_sequence(self):
        """Print the PKD sequence for inspection"""
        print("Periodic Kick Drive Sequence:")
        print("="*30)
        for i, (velocity_kick, interval) in enumerate(self.kick_sequence):
            print(f"Kick {i+1}: Δv=({velocity_kick[0]:.3f}, {velocity_kick[1]:.3f}), Interval={interval} steps")
        print()


class DampedOscillator2D:
    def __init__(self, k=1.0, c=0.2, m=1.0):
        """
        Initialize 2D damped harmonic oscillator parameters
        
        Parameters:
        - k: spring constant (stiffness)
        - c: damping coefficient 
        - m: mass
        """
        self.k = k  # Spring constant
        self.c = c  # Damping coefficient
        self.m = m  # Mass
        self.equilibrium = np.array([1.0, 1.0])  # Equilibrium position (1, 1)
        
    def system_dynamics(self, state, t):
        """
        Define the system dynamics for the 2D damped oscillator
        
        State vector: [x, vx, y, vy]
        
        Equations of motion:
        x' = vx
        vx' = -(c/m)*vx - (k/m)*(x - x0)
        y' = vy
        vy' = -(c/m)*vy - (k/m)*(y - y0)
        """
        x, vx, y, vy = state
        
        # Dynamics for x direction
        dxdt = vx
        dvxdt = -(self.c/self.m) * vx - (self.k/self.m) * (x - self.equilibrium[0])
        
        # Dynamics for y direction
        dydt = vy
        dvydt = -(self.c/self.m) * vy - (self.k/self.m) * (y - self.equilibrium[1])
        
        return [dxdt, dvxdt, dydt, dvydt]
    
    def simulate(self, initial_state, num_steps=1000, dt=0.01):
        """
        Simulate the 2D damped oscillator system
        
        Parameters:
        - initial_state: [x0, vx0, y0, vy0] initial conditions
        - num_steps: number of simulation steps
        - dt: time step
        
        Returns:
        - t: time array
        - trajectory: state trajectory [x, vx, y, vy] over time
        """
        t = np.linspace(0, num_steps * dt, num_steps)
        trajectory = odeint(self.system_dynamics, initial_state, t)
        
        return t, trajectory
    
    def simulate_with_pkd(self, initial_state, pkd, num_steps=3000, dt=0.01, pkd_start_ratio=0.3):
        """
        Simulate the 2D damped oscillator system with Periodic Kick Drive
        
        Parameters:
        - initial_state: [x0, vx0, y0, vy0] initial conditions
        - pkd: PeriodicKickDrive instance
        - num_steps: number of simulation steps
        - dt: time step
        - pkd_start_ratio: fraction of simulation where PKD starts (0.3 = 30%)
        
        Returns:
        - t: time array
        - trajectory: state trajectory [x, vx, y, vy] over time
        - kick_times: list of times when kicks were applied
        - kick_locations: list of positions where kicks were applied
        """
        # Initialize arrays to store results
        trajectory = np.zeros((num_steps, 4))
        t = np.zeros(num_steps)
        kick_times = []
        kick_locations = []
        
        # Set initial conditions
        current_state = np.array(initial_state)
        trajectory[0] = current_state
        t[0] = 0
        
        # Determine when PKD starts (random time after pkd_start_ratio of simulation)
        pkd_start_step = int(pkd_start_ratio * num_steps) + random.randint(0, int(0.2 * num_steps))
        pkd_start_step = min(pkd_start_step, num_steps - 100)  # Ensure we have room for kicks
        
        # Reset PKD sequence
        pkd.reset()
        
        # PKD state variables
        pkd_active = False
        steps_until_next_kick = 0
        current_kick_data = None
        
        print(f"PKD will start at step {pkd_start_step} (time {pkd_start_step * dt:.2f}s)")
        
        # Main simulation loop
        for step in range(1, num_steps):
            t[step] = step * dt
            
            # Check if we should activate PKD
            if step == pkd_start_step:
                pkd_active = True
                current_kick_data = pkd.get_current_kick()
                velocity_kick, interval = current_kick_data
                
                # Apply the first kick
                current_state[1] += velocity_kick[0]  # vx
                current_state[3] += velocity_kick[1]  # vy
                
                # Record kick
                kick_times.append(t[step])
                kick_locations.append([current_state[0], current_state[2]])
                
                # Set steps until next kick
                steps_until_next_kick = interval
                print(f"Applied kick at step {step}: Δv=({velocity_kick[0]:.3f}, {velocity_kick[1]:.3f})")
            
            # Check if we need to apply a kick
            elif pkd_active and steps_until_next_kick <= 0:
                current_kick_data = pkd.get_current_kick()
                velocity_kick, interval = current_kick_data
                
                # Apply the kick
                current_state[1] += velocity_kick[0]  # vx
                current_state[3] += velocity_kick[1]  # vy
                
                # Record kick
                kick_times.append(t[step])
                kick_locations.append([current_state[0], current_state[2]])
                
                # Set steps until next kick
                steps_until_next_kick = interval
                print(f"Applied kick at step {step}: Δv=({velocity_kick[0]:.3f}, {velocity_kick[1]:.3f})")
            
            # Integrate one time step using Euler's method for more control
            derivatives = self.system_dynamics(current_state, t[step])
            current_state = current_state + np.array(derivatives) * dt
            
            # Store the result
            trajectory[step] = current_state
            
            # Decrement kick counter
            if pkd_active:
                steps_until_next_kick -= 1
        
        print(f"PKD simulation completed. Applied {len(kick_times)} kicks total.")
        
        return t, trajectory, kick_times, kick_locations
    
    def generate_random_initial_conditions(self, position_range=3.0, velocity_range=2.0):
        """
        Generate random initial conditions within reasonable ranges
        
        Parameters:
        - position_range: range for initial positions around equilibrium
        - velocity_range: range for initial velocities
        
        Returns:
        - initial_state: [x0, vx0, y0, vy0]
        """
        x0 = self.equilibrium[0] + random.uniform(-position_range, position_range)
        y0 = self.equilibrium[1] + random.uniform(-position_range, position_range) 
        vx0 = random.uniform(-velocity_range, velocity_range)
        vy0 = random.uniform(-velocity_range, velocity_range)
        
        return [x0, vx0, y0, vy0]
    
    def compute_vector_field(self, x_range, y_range, grid_density=15):
        """
        Compute the vector field for the 2D damped oscillator
        
        Parameters:
        - x_range: tuple (x_min, x_max) for x-axis range
        - y_range: tuple (y_min, y_max) for y-axis range  
        - grid_density: number of grid points along each axis
        
        Returns:
        - X, Y: meshgrid coordinates
        - U, V: velocity components at each grid point
        """
        x = np.linspace(x_range[0], x_range[1], grid_density)
        y = np.linspace(y_range[0], y_range[1], grid_density)
        X, Y = np.meshgrid(x, y)
        
        # For position space, the vector field shows the velocity directions
        # At each position (x,y), we compute what the acceleration would be
        # assuming zero initial velocity (this gives us the force field direction)
        
        # Acceleration = -(k/m) * (position - equilibrium)
        U = -(self.k/self.m) * (X - self.equilibrium[0])  # acceleration in x
        V = -(self.k/self.m) * (Y - self.equilibrium[1])  # acceleration in y
        
        return X, Y, U, V
    
    def visualize_trajectory(self, trajectory, save_path=None, show_vector_field=True, kick_locations=None, kick_times=None, t=None):
        """
        Visualize the 2D trajectory in position space with optional vector field and kick locations
        Only shows the last 5 PKD cycles (at most 25 kicks)
        
        Parameters:
        - trajectory: state trajectory from simulation
        - save_path: optional path to save the plot
        - show_vector_field: whether to overlay the vector field
        - kick_locations: list of [x, y] positions where kicks were applied
        - kick_times: list of times when kicks were applied
        - t: time array from simulation
        """
        # Determine the start index for the last 5 PKD cycles (5 kicks per cycle)
        max_cycles = 5
        max_kicks = max_cycles * 5
        
        if kick_times is not None and len(kick_times) > max_kicks:
            # Find the time when the last 5 cycles started
            start_time = kick_times[-max_kicks]
            # Find the corresponding index in the trajectory
            if t is not None:
                start_idx = np.searchsorted(t, start_time)
            else:
                start_idx = 0
        else:
            start_idx = 0
        
        # Use only the trajectory from the start of the last 5 cycles
        x_pos = trajectory[start_idx:, 0]
        y_pos = trajectory[start_idx:, 2]
        
        # Filter kick locations to only show those from the last 5 cycles
        filtered_kick_locations = None
        if kick_locations is not None and kick_times is not None and len(kick_times) > max_kicks:
            filtered_kick_locations = kick_locations[-max_kicks:]
        else:
            filtered_kick_locations = kick_locations
        
        plt.figure(figsize=(10, 8))
        
        # Add vector field if requested
        if show_vector_field:
            # Determine plot range based on trajectory data with extra padding
            x_min, x_max = min(x_pos.min(), self.equilibrium[0]-2), max(x_pos.max(), self.equilibrium[0]+2)
            y_min, y_max = min(y_pos.min(), self.equilibrium[1]-2), max(y_pos.max(), self.equilibrium[1]+2)
            
            # Add generous padding to cover the entire plot area
            x_range = (x_min - 1.5, x_max + 1.5)
            y_range = (y_min - 1.5, y_max + 1.5)
            
            # Compute vector field with higher density for better coverage
            X, Y, U, V = self.compute_vector_field(x_range, y_range, grid_density=15)
            
            # Normalize the vectors to have uniform arrow lengths
            magnitude = np.sqrt(U**2 + V**2)
            # Avoid division by zero
            magnitude = np.where(magnitude == 0, 1, magnitude)
            U_norm = U / magnitude
            V_norm = V / magnitude
            
            # Plot normalized vector field with semi-transparent gray arrows (3/4 length)
            plt.quiver(X, Y, U_norm, V_norm, alpha=0.4, color='gray', scale=33, width=0.002,
                      headwidth=3, headlength=4, label='Force Field (Normalized)')
        
        # Plot the trajectory as individual segments for cumulative transparency effect
        for i in range(len(x_pos) - 1):
            if i == 0:  # Add label only for the first segment
                plt.plot([x_pos[i], x_pos[i+1]], [y_pos[i], y_pos[i+1]], 
                        'b-', linewidth=2, alpha=0.1, zorder=3, label='Trajectory (Last 5 Cycles)')
            else:
                plt.plot([x_pos[i], x_pos[i+1]], [y_pos[i], y_pos[i+1]], 
                        'b-', linewidth=2, alpha=0.1, zorder=3)
        
        # Mark kick locations if provided (filtered to last 5 cycles)
        if filtered_kick_locations is not None and len(filtered_kick_locations) > 0:
            kick_x = [loc[0] for loc in filtered_kick_locations]
            kick_y = [loc[1] for loc in filtered_kick_locations]
            plt.scatter(kick_x, kick_y, c='orange', s=80, marker='x', linewidth=3,
                       label=f'PKD Kicks ({len(filtered_kick_locations)})', zorder=5)
        
        # Mark starting point
        plt.plot(x_pos[0], y_pos[0], 'go', markersize=8, label='Start', zorder=4)
        
        # Mark ending point
        plt.plot(x_pos[-1], y_pos[-1], 'ro', markersize=8, label='End', zorder=4)
        
        # Mark equilibrium point
        plt.plot(self.equilibrium[0], self.equilibrium[1], 'k*', markersize=12, 
                label=f'Equilibrium ({self.equilibrium[0]}, {self.equilibrium[1]})', zorder=4)
        
        plt.xlabel('X Position')
        plt.ylabel('Y Position')
        title = '2D Damped Harmonic Oscillator Trajectory with Normalized Force Field'
        if kick_locations is not None and len(kick_locations) > 0:
            title += f' (PKD: {len(kick_locations)} kicks)'
        plt.title(title)
        plt.grid(True, alpha=0.3)
        plt.legend()
        plt.axis('equal')
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        
        plt.show()
    
    def save_trajectory_svg(self, trajectory, kick_locations=None, kick_times=None, t=None, 
                           svg_path="./logs/oscillator_trajectory.svg", show_vector_field=True):
        """
        Save the 2D trajectory visualization as SVG (last 5 PKD cycles only)
        
        Parameters:
        - trajectory: state trajectory from simulation
        - kick_locations: list of [x, y] positions where kicks were applied
        - kick_times: list of times when kicks were applied
        - t: time array from simulation
        - svg_path: path to save the SVG file
        - show_vector_field: whether to include vector field
        """
        # Determine the start index for the last 5 PKD cycles (5 kicks per cycle)
        max_cycles = 5
        max_kicks = max_cycles * 5
        
        if kick_times is not None and len(kick_times) > max_kicks:
            # Find the time when the last 5 cycles started
            start_time = kick_times[-max_kicks]
            # Find the corresponding index in the trajectory
            if t is not None:
                start_idx = np.searchsorted(t, start_time)
            else:
                start_idx = 0
        else:
            start_idx = 0
        
        # Use only the trajectory from the start of the last 5 cycles
        x_pos = trajectory[start_idx:, 0]
        y_pos = trajectory[start_idx:, 2]
        
        # Filter kick locations to only show those from the last 5 cycles
        filtered_kick_locations = None
        if kick_locations is not None and kick_times is not None and len(kick_times) > max_kicks:
            filtered_kick_locations = kick_locations[-max_kicks:]
        else:
            filtered_kick_locations = kick_locations
        
        # Determine tight plot range based on actual trajectory bounds
        x_min, x_max = x_pos.min(), x_pos.max()
        y_min, y_max = y_pos.min(), y_pos.max()
        
        # Include equilibrium point in bounds
        x_min = min(x_min, self.equilibrium[0])
        x_max = max(x_max, self.equilibrium[0])
        y_min = min(y_min, self.equilibrium[1])
        y_max = max(y_max, self.equilibrium[1])
        
        # Calculate data range
        x_data_range = x_max - x_min
        y_data_range = y_max - y_min
        
        # Set margin in data coordinates (equivalent to 100 pixels)
        margin_pixels = 100
        base_svg_size = 600  # Base size to calculate margin ratio
        
        # Calculate margin in data coordinates
        x_margin = (x_data_range * margin_pixels) / (base_svg_size - 2 * margin_pixels)
        y_margin = (y_data_range * margin_pixels) / (base_svg_size - 2 * margin_pixels)
        
        # Apply margins
        x_range = (x_min - x_margin, x_max + x_margin)
        y_range = (y_min - y_margin, y_max + y_margin)
        
        # Calculate SVG dimensions based on data aspect ratio
        data_width = x_range[1] - x_range[0]
        data_height = y_range[1] - y_range[0]
        aspect_ratio = data_width / data_height
        
        # Set SVG dimensions to maintain aspect ratio
        if aspect_ratio > 1:
            svg_width = 800
            svg_height = int(800 / aspect_ratio)
        else:
            svg_height = 600
            svg_width = int(600 * aspect_ratio)
        
        # Create coordinate transformation
        x_scale = svg_width / data_width
        y_scale = svg_height / data_height
        
        def transform_coords(x, y):
            svg_x = (x - x_range[0]) * x_scale
            svg_y = svg_height - (y - y_range[0]) * y_scale  # Flip Y axis
            return svg_x, svg_y
        
        # Create SVG drawing
        dwg = svgwrite.Drawing(svg_path, size=(svg_width, svg_height))
        
        # Add background
        dwg.add(dwg.rect(insert=(0, 0), size=(svg_width, svg_height), fill='white'))
        
        # Add vector field if requested
        if show_vector_field:
            # Compute vector field
            X, Y, U, V = self.compute_vector_field(x_range, y_range, grid_density=15)
            
            # Normalize vectors
            magnitude = np.sqrt(U**2 + V**2)
            magnitude = np.where(magnitude == 0, 1, magnitude)
            U_norm = U / magnitude
            V_norm = V / magnitude
            
            # Draw vector field arrows
            arrow_scale = 15  # Scale factor for arrow length
            for i in range(X.shape[0]):
                for j in range(X.shape[1]):
                    x_start, y_start = transform_coords(X[i, j], Y[i, j])
                    x_end = x_start + U_norm[i, j] * arrow_scale
                    y_end = y_start - V_norm[i, j] * arrow_scale  # Flip Y component
                    
                    # Draw arrow line
                    dwg.add(dwg.line(start=(x_start, y_start), end=(x_end, y_end),
                                   stroke='gray', stroke_width=1, opacity=0.4))
                    
                    # Draw arrowhead
                    arrow_size = 3
                    angle = np.arctan2(-V_norm[i, j], U_norm[i, j])  # Flip Y for SVG
                    
                    # Arrowhead points
                    p1_x = x_end - arrow_size * np.cos(angle - np.pi/6)
                    p1_y = y_end - arrow_size * np.sin(angle - np.pi/6)
                    p2_x = x_end - arrow_size * np.cos(angle + np.pi/6)
                    p2_y = y_end - arrow_size * np.sin(angle + np.pi/6)
                    
                    dwg.add(dwg.polygon(points=[(x_end, y_end), (p1_x, p1_y), (p2_x, p2_y)],
                                      fill='gray', opacity=0.4))
        
        # Draw trajectory as individual segments with full opacity
        for i in range(len(x_pos) - 1):
            x1, y1 = transform_coords(x_pos[i], y_pos[i])
            x2, y2 = transform_coords(x_pos[i+1], y_pos[i+1])
            
            dwg.add(dwg.line(start=(x1, y1), end=(x2, y2),
                           stroke='blue', stroke_width=2, opacity=1.0))
        
        # Mark kick locations with small red circles
        if filtered_kick_locations is not None and len(filtered_kick_locations) > 0:
            for kick_loc in filtered_kick_locations:
                kick_x, kick_y = transform_coords(kick_loc[0], kick_loc[1])
                # Draw small red circle for kicks (radius = stroke_width/2 = 1)
                dwg.add(dwg.circle(center=(kick_x, kick_y), r=1, fill='red'))
        
        # Mark equilibrium point with small black circle
        eq_x, eq_y = transform_coords(self.equilibrium[0], self.equilibrium[1])
        dwg.add(dwg.circle(center=(eq_x, eq_y), r=3, fill='black'))
        
        # Save SVG
        dwg.save()
        print(f"SVG trajectory saved to {svg_path}")
    
    def compute_phase_space_vector_field(self, pos_range, vel_range, direction='x', grid_density=12):
        """
        Compute the vector field in phase space for one direction
        
        Parameters:
        - pos_range: tuple (pos_min, pos_max) for position axis range
        - vel_range: tuple (vel_min, vel_max) for velocity axis range
        - direction: 'x' or 'y' to specify which direction
        - grid_density: number of grid points along each axis
        
        Returns:
        - POS, VEL: meshgrid coordinates
        - U, V: phase space vector components (dpos/dt, dvel/dt)
        """
        pos = np.linspace(pos_range[0], pos_range[1], grid_density)
        vel = np.linspace(vel_range[0], vel_range[1], grid_density)
        POS, VEL = np.meshgrid(pos, vel)
        
        # Phase space dynamics:
        # dpos/dt = vel
        # dvel/dt = -(c/m)*vel - (k/m)*(pos - equilibrium)
        
        equilibrium_pos = self.equilibrium[0] if direction == 'x' else self.equilibrium[1]
        
        U = VEL  # dpos/dt = velocity
        V = -(self.c/self.m) * VEL - (self.k/self.m) * (POS - equilibrium_pos)  # dvel/dt
        
        return POS, VEL, U, V
    
    def visualize_phase_space(self, trajectory, save_path=None, show_vector_field=True, kick_times=None, t=None):
        """
        Visualize phase space trajectories for both x and y directions with optional vector fields
        Only shows the last 5 PKD cycles (at most 25 kicks)
        
        Parameters:
        - trajectory: state trajectory from simulation
        - save_path: optional path to save the plot
        - show_vector_field: whether to overlay the vector fields
        - kick_times: list of times when kicks were applied
        - t: time array from simulation
        """
        # Determine the start index for the last 5 PKD cycles (5 kicks per cycle)
        max_cycles = 5
        max_kicks = max_cycles * 5
        
        if kick_times is not None and len(kick_times) > max_kicks:
            # Find the time when the last 5 cycles started
            start_time = kick_times[-max_kicks]
            # Find the corresponding index in the trajectory
            if t is not None:
                start_idx = np.searchsorted(t, start_time)
            else:
                start_idx = 0
        else:
            start_idx = 0
        
        # Use only the trajectory from the start of the last 5 cycles
        x_pos = trajectory[start_idx:, 0]
        x_vel = trajectory[start_idx:, 1]
        y_pos = trajectory[start_idx:, 2]
        y_vel = trajectory[start_idx:, 3]
        
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
        
        # X direction phase space
        if show_vector_field:
            # Determine ranges for vector field
            x_pos_range = (min(x_pos.min(), self.equilibrium[0]-0.5), 
                          max(x_pos.max(), self.equilibrium[0]+0.5))
            x_vel_range = (x_vel.min()-0.2, x_vel.max()+0.2)
            
            # Compute and plot vector field
            POS_X, VEL_X, U_X, V_X = self.compute_phase_space_vector_field(
                x_pos_range, x_vel_range, 'x', grid_density=10)
            
            ax1.quiver(POS_X, VEL_X, U_X, V_X, alpha=0.4, color='gray', scale=20, width=0.003,
                      headwidth=3, headlength=4)
        
        # Plot x-direction trajectory as individual segments for cumulative transparency
        for i in range(len(x_pos) - 1):
            ax1.plot([x_pos[i], x_pos[i+1]], [x_vel[i], x_vel[i+1]], 
                    'b-', linewidth=2, alpha=0.1, zorder=3)
        ax1.plot(x_pos[0], x_vel[0], 'go', markersize=8, label='Start', zorder=4)
        ax1.plot(x_pos[-1], x_vel[-1], 'ro', markersize=8, label='End', zorder=4)
        ax1.plot(self.equilibrium[0], 0, 'k*', markersize=12, label='Equilibrium', zorder=4)
        ax1.set_xlabel('X Position')
        ax1.set_ylabel('X Velocity')
        ax1.set_title('X Direction Phase Space with Vector Field')
        ax1.grid(True, alpha=0.3)
        ax1.legend()
        
        # Y direction phase space  
        if show_vector_field:
            # Determine ranges for vector field
            y_pos_range = (min(y_pos.min(), self.equilibrium[1]-0.5),
                          max(y_pos.max(), self.equilibrium[1]+0.5))
            y_vel_range = (y_vel.min()-0.2, y_vel.max()+0.2)
            
            # Compute and plot vector field
            POS_Y, VEL_Y, U_Y, V_Y = self.compute_phase_space_vector_field(
                y_pos_range, y_vel_range, 'y', grid_density=10)
            
            ax2.quiver(POS_Y, VEL_Y, U_Y, V_Y, alpha=0.4, color='gray', scale=20, width=0.003,
                      headwidth=3, headlength=4)
        
        # Plot y-direction trajectory as individual segments for cumulative transparency
        for i in range(len(y_pos) - 1):
            ax2.plot([y_pos[i], y_pos[i+1]], [y_vel[i], y_vel[i+1]], 
                    'r-', linewidth=2, alpha=0.1, zorder=3)
        ax2.plot(y_pos[0], y_vel[0], 'go', markersize=8, label='Start', zorder=4)
        ax2.plot(y_pos[-1], y_vel[-1], 'ro', markersize=8, label='End', zorder=4)
        ax2.plot(self.equilibrium[1], 0, 'k*', markersize=12, label='Equilibrium', zorder=4)
        ax2.set_xlabel('Y Position')
        ax2.set_ylabel('Y Velocity')
        ax2.set_title('Y Direction Phase Space with Vector Field')
        ax2.grid(True, alpha=0.3)
        ax2.legend()
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        
        plt.show()
    
    def print_system_info(self):
        """Print system parameters and information"""
        print("2D Damped Harmonic Oscillator System")
        print("="*40)
        print(f"Mass (m): {self.m}")
        print(f"Spring constant (k): {self.k}")
        print(f"Damping coefficient (c): {self.c}")
        print(f"Equilibrium position: ({self.equilibrium[0]}, {self.equilibrium[1]})")
        print(f"Natural frequency: {np.sqrt(self.k/self.m):.3f} rad/s")
        print(f"Damping ratio: {self.c/(2*np.sqrt(self.k*self.m)):.3f}")
        print()


def main():
    """Main function to run the 2D damped oscillator simulation with PKD"""
    
    # Initialize the oscillator with reasonable parameters
    # k=1.5 (moderate stiffness), c=0.3 (light damping), m=1.0
    oscillator = DampedOscillator2D(k=1.5, c=0.3, m=1.0)
    
    # Initialize Periodic Kick Drive (PKD) - shorter intervals for more cycles
    pkd = PeriodicKickDrive(n_kicks=5, velocity_range=1.5, interval_range=(30, 80))
    
    # Print system information
    oscillator.print_system_info()
    pkd.print_sequence()
    
    # Generate random initial conditions
    initial_state = oscillator.generate_random_initial_conditions(
        position_range=2.5, velocity_range=3.0
    )
    
    print(f"Initial conditions:")
    print(f"  Position: ({initial_state[0]:.3f}, {initial_state[2]:.3f})")
    print(f"  Velocity: ({initial_state[1]:.3f}, {initial_state[3]:.3f})")
    print()
    
    # Run PKD simulation for 10000 steps
    print("Running PKD simulation for 10000 steps...")
    t, trajectory, kick_times, kick_locations = oscillator.simulate_with_pkd(
        initial_state, pkd, num_steps=10000, dt=0.01, pkd_start_ratio=0.3
    )
    
    print(f"PKD simulation completed. Total time: {t[-1]:.2f} seconds")
    print(f"Final position: ({trajectory[-1, 0]:.3f}, {trajectory[-1, 2]:.3f})")
    print(f"Final velocity: ({trajectory[-1, 1]:.3f}, {trajectory[-1, 3]:.3f})")
    print(f"Kick times: {[f'{kt:.2f}s' for kt in kick_times]}")
    print()
    
    # Ensure logs directory exists
    os.makedirs("./logs", exist_ok=True)
    
    # Generate SVG export with timestamp
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    svg_filename = f"./logs/oscillator_trajectory_{timestamp}.svg"
    
    print("Exporting SVG...")
    oscillator.save_trajectory_svg(trajectory, kick_locations=kick_locations, 
                                  kick_times=kick_times, t=t, svg_path=svg_filename)
    
    # Visualize the 2D trajectory with kick locations (last 5 cycles only)
    print("Generating visualization...")
    oscillator.visualize_trajectory(trajectory, kick_locations=kick_locations, kick_times=kick_times, t=t)
    
    # Also show phase space plots (last 5 cycles only)
    oscillator.visualize_phase_space(trajectory, kick_times=kick_times, t=t)
    
    # Calculate and display some statistics
    x_positions = trajectory[:, 0]
    y_positions = trajectory[:, 2]
    distances = np.sqrt((x_positions - oscillator.equilibrium[0])**2 + 
                       (y_positions - oscillator.equilibrium[1])**2)
    
    print("Trajectory Statistics:")
    print(f"  Maximum distance from equilibrium: {np.max(distances):.3f}")
    print(f"  Final distance from equilibrium: {distances[-1]:.3f}")
    print(f"  Average distance from equilibrium: {np.mean(distances):.3f}")
    print(f"  PKD kicks applied: {len(kick_times)} at times {[f'{kt:.1f}s' for kt in kick_times]}")
    
    return oscillator, trajectory, t, pkd, kick_times, kick_locations


if __name__ == "__main__":
    # Set random seed for reproducibility (optional)
    random.seed()
    np.random.seed()
    
    # Run the main simulation
    oscillator, trajectory, time, pkd, kick_times, kick_locations = main()
